package com.aiwiown.snackmq.api.impl;

import com.aiwiown.snackmq.api.SendResult;
import com.aiwiown.snackmq.api.config.ProducerConfig;
import com.aiwiown.snackmq.common.dto.TopicInfo;
import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.protocol.TopicPartition;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 分区级批处理管理器
 * 它将消息收集到按分区组织的批次中，并根据大小或时间判断批次是否就绪。
 * 这个类是线程安全的。
 */
@Slf4j
public class PartitionBatchManager {

    private final ProducerConfig producerConfig;
    private final Map<TopicPartition, MessageBatch> partitionBatches = new ConcurrentHashMap<>();
    private final AtomicLong totalMessagesBatched = new AtomicLong(0);
    // 【优化】: 添加一个函数式接口，用于从外部获取主题信息
    private final Function<String, TopicInfo> topicInfoProvider;
    /**
     * 一个包装器，将消息与其对应的 Future 关联起来。
     */
    @Getter
    @AllArgsConstructor
    public static class MessageFutureWrapper {
        private final Message message;
        private final CompletableFuture<SendResult> future;
    }

    /**
     * 单个分区的消息批次，现在存储 MessageFutureWrapper。
     */
    @Getter
    public static class MessageBatch {
        private final TopicPartition partition;
        private final List<MessageFutureWrapper> messageWrappers = new ArrayList<>();
        private final long createTime;

        public MessageBatch(TopicPartition partition) {
            this.partition = partition;
            this.createTime = System.currentTimeMillis();
        }

        public void add(MessageFutureWrapper wrapper) {
            this.messageWrappers.add(wrapper);
        }

        public boolean isFull(int batchSize) {
            return this.messageWrappers.size() >= batchSize;
        }

        public boolean isExpired(long lingerMs) {
            // 只有当 lingerMs > 0 时，才检查超时
            return lingerMs > 0 && (System.currentTimeMillis() - createTime) >= lingerMs;
        }

        public int size() {
            return messageWrappers.size();
        }

        public List<Message> getMessages() {
            return messageWrappers.stream()
                    .map(MessageFutureWrapper::getMessage)
                    .collect(Collectors.toList());
        }
    }

    public PartitionBatchManager(ProducerConfig producerConfig, Function<String, TopicInfo> topicInfoProvider) {
        this.producerConfig = producerConfig;
        this.topicInfoProvider = topicInfoProvider;
    }

    /**
     * 添加一条消息到相应的批次中。
     * 如果添加后批次已满，则将该批次从管理器中移除并返回，以便立即发送。
     *
     * @param wrapper 包含消息和其 Future 的包装器
     * @return 如果批次已满，则返回该批次；否则返回 null。
     */
    public synchronized MessageBatch add(MessageFutureWrapper wrapper) {
        TopicPartition partition = calculatePartition(wrapper.getMessage());
        MessageBatch batch = partitionBatches.computeIfAbsent(partition, MessageBatch::new);

        batch.add(wrapper);
        totalMessagesBatched.incrementAndGet();

        if (batch.isFull(producerConfig.getBatchSize())) {
            // 批次已满，从 map 中移除并返回
            return partitionBatches.remove(partition);
        }

        return null;
    }

    /**
     * 获取所有因超时而就绪的批次。
     * 此方法会从管理器中移除这些批次。
     *
     * @return 所有超时的批次列表。
     */
    public synchronized List<MessageBatch> getExpiredBatches() {
        List<MessageBatch> expiredBatches = new ArrayList<>();
        // 使用迭代器安全地移除元素
        partitionBatches.entrySet().removeIf(entry -> {
            if (entry.getValue().isExpired(producerConfig.getLingerMs())) {
                expiredBatches.add(entry.getValue());
                return true; // 从 map 中移除
            }
            return false;
        });
        return expiredBatches;
    }

    /**
     * 关闭管理器，并返回所有剩余的、未发送的批次。
     *
     * @return 所有剩余的批次列表。
     */
    public synchronized List<MessageBatch> close() {
        List<MessageBatch> remainingBatches = new ArrayList<>(partitionBatches.values());
        partitionBatches.clear();
        log.info("PartitionBatchManager closed. Drained {} remaining batches.", remainingBatches.size());
        return remainingBatches;
    }

    /**
     * 计算消息应该发送到哪个分区。
     */
    private TopicPartition calculatePartition(Message message) {
        String topic = message.getTopic();
        // 【最终修复】: 修复了 Exchange 路由的逻辑。
        // 旧的实现错误地检查了 "exchange" 属性，而实际上应该检查 "exchangeName"。
        String exchangeName = message.getProperty("exchangeName");
        String key = message.getKey();
        
        // 【修复】优先检查是否是发送到交换机的消息
        if (exchangeName != null && !exchangeName.isEmpty()) {
            // 模式二：通过 Exchange 发送
            // 对于发往 Exchange 的消息，我们使用一个特殊的虚拟主题来分组批次
            // 这样状态机就能正确识别这是交换机消息而不是普通消息
            return new TopicPartition("__EXCHANGE_ROUTING__", 0);
        }
        
        // 模式一：直接发送到 Topic
        if (topic != null && !topic.isEmpty()) {
            int partitionCount = 1;
            TopicInfo topicInfo = topicInfoProvider.apply(topic);
            if (topicInfo != null && topicInfo.getPartitionCount() > 0) {
                partitionCount = topicInfo.getPartitionCount();
            } else {
                log.warn("无法获取主题 '{}' 的元数据，将使用默认分区数 1。这可能导致消息路由不正确。", topic);
            }
            int partitionId = 0;
            if (key != null && !key.isEmpty() && partitionCount > 1) {
                partitionId = (key.hashCode() & 0x7FFFFFFF) % partitionCount;
            }
            return new TopicPartition(topic, partitionId);
        }

        // 异常情况：如果消息既没有 Topic 也没有 Exchange，这是一个错误。
        // 我们创建一个特殊的 TopicPartition 来隔离这些错误消息，防止它们污染正常批次。
        log.error("消息既没有 'topic' 也没有 'exchangeName' 属性，无法确定批处理分组: {}", message);
        return new TopicPartition("__INVALID_MESSAGE__", -1);
    }

    /**
     * 【新增】根据消息的 key 计算其应被发送到的分区ID。
     *
     * @param topic 主题名称
     * @param key   消息的 key，用于分区计算
     * @return 分区ID
     */
    public int getPartition(String topic, String key) {
        TopicInfo topicInfo = topicInfoProvider.apply(topic);
        int partitionCount = 1;
        if (topicInfo != null && topicInfo.getPartitionCount() > 0) {
            partitionCount = topicInfo.getPartitionCount();
        } else {
            log.warn("无法获取主题 '{}' 的元数据，将使用默认分区数 1。这可能导致消息路由不正确。", topic);
        }
        if (key != null && !key.isEmpty() && partitionCount > 1) {
            return (key.hashCode() & 0x7FFFFFFF) % partitionCount;
        }

        return 0; // 如果没有 key 或只有一个分区，则默认为分区 0
    }

    /**
     * 仅用于测试：获取所有分区批次的只读视图。
     */
    Map<TopicPartition, MessageBatch> getPartitionBatches() {
        return Collections.unmodifiableMap(partitionBatches);
    }
}